package com.tyros.common.secure.asymmetrics;

import com.tyros.common.secure.CryptoException;
import com.tyros.common.secure.GlobalProvider;
import com.tyros.common.secure.codec.Codec;
import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1Object;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERBitString;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;

import javax.crypto.Cipher;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.StringWriter;
import java.nio.charset.StandardCharsets;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.interfaces.RSAKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

/**
 * description: RSA非对称加密算法
 *
 * @author zsy
 */
public class RSA {
    public static final String CHARSET = "UTF-8";
    public static final String ALGORITHM = "RSA";
    public static final String ALGORITHM_PATTERN_PADDING = ALGORITHM + "/ECB/OAEPPadding";

    /**
     * 生成公钥和私钥
     *
     * @param keySize 密钥长度
     * @return 公钥密钥对
     */
    public static Map<String, String> createKeys(int keySize) {
        //为RSA算法创建一个KeyPairGenerator对象
        KeyPairGenerator kpg;
        try {
            kpg = KeyPairGenerator.getInstance(ALGORITHM);
        } catch (NoSuchAlgorithmException e) {
            throw new CryptoException("No such algorithm:" + ALGORITHM, e);
        }
        //初始化KeyPairGenerator对象,密钥长度
        kpg.initialize(keySize);
        //生成密匙对
        KeyPair keyPair = kpg.generateKeyPair();
        //得到公钥
        Key publicKey = keyPair.getPublic();
        String publicKeyStr = Codec.BASE64.encode(publicKey.getEncoded());
        //得到私钥
        Key privateKey = keyPair.getPrivate();
        String privateKeyStr = Codec.BASE64.encode(privateKey.getEncoded());
        Map<String, String> keyPairMap = new HashMap<String, String>(2);
        keyPairMap.put("publicKey", publicKeyStr);
        keyPairMap.put("privateKey", privateKeyStr);
        return keyPairMap;
    }

    /**
     * 通过X509编码的Key指令获得公钥对象
     *
     * @param publicKey 密钥字符串（经过base64编码）
     * @throws CryptoException 加密异常
     */
    public static RSAPublicKey getPublicKey(String publicKey) throws CryptoException {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
            X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(Codec.BASE64.decode(publicKey));
            return (RSAPublicKey) keyFactory.generatePublic(x509KeySpec);
        } catch (Exception e) {
            throw new CryptoException("", e);
        }
    }

    /**
     * 通过PKCS#8编码的Key指令获得私钥对象
     *
     * @param privateKey 密钥字符串（经过base64编码）
     * @throws CryptoException 加密异常
     */
    public static RSAPrivateKey getPrivateKey(String privateKey) throws CryptoException {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance(ALGORITHM);
            PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(Codec.BASE64.decode(privateKey));
            return (RSAPrivateKey) keyFactory.generatePrivate(pkcs8KeySpec);
        } catch (Exception e) {
            throw new CryptoException("", e);
        }
    }

    /**
     * 公钥加密
     *
     * @param rawData 原始内容
     * @param key     密钥
     * @return 密文
     * @throws CryptoException 加密异常
     */
    public static String encrypt(byte[] rawData, Key key) throws CryptoException {
        try {
            Cipher cipher = Cipher.getInstance(ALGORITHM_PATTERN_PADDING, GlobalProvider.getDefaultProvider());
            cipher.init(Cipher.ENCRYPT_MODE, key);
            return Codec.BASE64.encode(rsaSplitCodec(cipher, Cipher.ENCRYPT_MODE, rawData, ((RSAKey) key).getModulus().bitLength()));
        } catch (Exception e) {
            throw new CryptoException("Encrypt exception by public key ", e);
        }
    }

    /**
     * 公钥加密
     *
     * @param rawData 原始内容
     * @param key     密钥
     * @param keyType 密钥类型
     * @return 密文
     */
    public static String encrypt(String rawData, String key, KeyType keyType) {
        return encrypt(rawData.getBytes(StandardCharsets.UTF_8), keyType.getFunction().apply(key));
    }

    public static String encrypt(byte[] rawData, String key, KeyType keyType) {
        return encrypt(rawData, keyType.getFunction().apply(key));
    }

    /**
     * 私钥解密
     *
     * @param encryptedData 密文数据
     * @param key           密钥
     * @return 解密后的明文
     */
    public static String decrypt(byte[] encryptedData, Key key) {
        try {
            Cipher cipher = Cipher.getInstance(ALGORITHM_PATTERN_PADDING, GlobalProvider.getDefaultProvider());
            cipher.init(Cipher.DECRYPT_MODE, key);
            return new String(rsaSplitCodec(cipher, Cipher.DECRYPT_MODE, Codec.BASE64.decode(encryptedData), ((RSAKey) key).getModulus().bitLength()), CHARSET);
        } catch (Exception e) {
            throw new CryptoException("Decrypt exception by private key ", e);
        }
    }

    /**
     * 私钥解密
     *
     * @param encryptedData 密文数据
     * @param key           密钥
     * @param keyType       密钥类型
     * @return 解密后的明文
     */
    public static String decrypt(String encryptedData, String key, KeyType keyType) {
        return decrypt(encryptedData.getBytes(StandardCharsets.UTF_8), keyType.getFunction().apply(key));
    }


    /**
     * 分段加解密
     *
     * @param cipher
     * @param opmode
     * @param datas
     * @param keySize
     * @return
     */
    private static byte[] rsaSplitCodec(Cipher cipher, int opmode, byte[] datas, int keySize) throws IOException {
        int maxBlock = 0;
        if (opmode == Cipher.DECRYPT_MODE) {
            maxBlock = keySize / 8;
        } else {
            maxBlock = keySize / 8 - 11;
        }
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        byte[] buff;
        int i = 0;
        try {
            while (datas.length > offSet) {
                if (datas.length - offSet > maxBlock) {
                    buff = cipher.doFinal(datas, offSet, maxBlock);
                } else {
                    buff = cipher.doFinal(datas, offSet, datas.length - offSet);
                }
                out.write(buff, 0, buff.length);
                i++;
                offSet = i * maxBlock;
            }
        } catch (Exception e) {
            throw new CryptoException("An exception occurred while " + opmode + " data", e);
        }
        byte[] resultData = out.toByteArray();
        out.close();
        return resultData;
    }


    /**
     * 公钥pkcs1转pkcs8
     *
     * @param pkcs1PublicKeyByte
     * @return pkcs8公钥
     */
    public static byte[] formatPublicKeyPKCS1ToPKCS8(byte[] pkcs1PublicKeyByte) {
        org.bouncycastle.asn1.pkcs.RSAPublicKey rsaPub = org.bouncycastle.asn1.pkcs.RSAPublicKey.getInstance(pkcs1PublicKeyByte);
        byte[] pkcs8Bytes = null;
        try {
            KeyFactory kf = KeyFactory.getInstance(ALGORITHM);
            PublicKey generatePublic = kf.generatePublic(new RSAPublicKeySpec(rsaPub.getModulus(), rsaPub.getPublicExponent()));
            pkcs8Bytes = generatePublic.getEncoded();
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new CryptoException("An exception occurred while parsing public key", e);
        }
        return pkcs8Bytes;
    }

    /**
     * 公钥pkcs8转pkcs1
     *
     * @param pkcs8PublicKeyByte
     * @return pkcs1公钥
     */
    public static byte[] formatPublicKeyPKCS8ToPKCS1(byte[] pkcs8PublicKeyByte) {
        ASN1Sequence asn1Encodables = ASN1Sequence.getInstance(pkcs8PublicKeyByte);
        ASN1Encodable asn1Encodable = asn1Encodables.getObjectAt(1);
        DERBitString derBitStringObject = DERBitString.getInstance(asn1Encodable);
        return derBitStringObject.getBytes();
    }

    /**
     * 私钥pkcs1转pkcs8
     *
     * @param pkcs1PrivateKeyByte
     * @return
     */
    public static byte[] formatPrivateKeyPKCS1ToPKCS8(byte[] pkcs1PrivateKeyByte) {
        AlgorithmIdentifier algorithmIdentifier = new AlgorithmIdentifier(PKCSObjectIdentifiers.pkcs8ShroudedKeyBag);
        //        org.bouncycastle.asn1.pkcs.RSAPrivateKey privateKey = org.bouncycastle.asn1.pkcs.RSAPrivateKey.getInstance(pkcs1PrivateKeyByte);
        //        ASN1Sequence privateKey = ASN1Sequence.getInstance(pkcs1PrivateKeyByte);
        try {
            ASN1Object asn1Object = ASN1ObjectIdentifier.fromByteArray(pkcs1PrivateKeyByte);
            PrivateKeyInfo privKeyInfo = new PrivateKeyInfo(algorithmIdentifier, asn1Object);
            return privKeyInfo.getEncoded();
        } catch (IOException e) {
            throw new CryptoException("An exception occurred while parsing private key", e);
        }
    }

    /**
     * 私钥pkcs8转pkcs1
     *
     * @param pksc8PrivateKeyByte
     * @return
     */
    public static byte[] formatPrivateKeyPKCS8ToPKCS1(byte[] pksc8PrivateKeyByte) {
        PrivateKeyInfo pki = PrivateKeyInfo.getInstance(pksc8PrivateKeyByte);
        byte[] pkcs1Bytes = null;
        try {
            ASN1Encodable parsePrivateKey = pki.parsePrivateKey();
            pkcs1Bytes = parsePrivateKey.toASN1Primitive().getEncoded();
        } catch (IOException e) {
            throw new CryptoException("An exception occurred while parsing private key", e);
        }
        return pkcs1Bytes;

    }

    public static String format2Pem(String type, byte[] privateKeyPKCS1) throws Exception {
        PemObject pemObject = new PemObject(type, privateKeyPKCS1);
        StringWriter stringWriter = new StringWriter();
        PemWriter pemWriter = new PemWriter(stringWriter);
        pemWriter.writeObject(pemObject);
        pemWriter.close();
        return stringWriter.toString();
    }

    public enum KeyType {
        /**
         * 私钥
         */
        PRIVATE_KEY(RSA::getPrivateKey),
        /**
         * 公钥
         */
        PUBLIC_KEY(RSA::getPublicKey);
        private final Function<String, Key> function;

        KeyType(Function<String, Key> function) {
            this.function = function;
        }

        public Function<String, Key> getFunction() {
            return function;
        }
    }
}
